/**
 * @file logistic_regression_main.cpp
 * @author Ryan Curtin
 *
 * Main executable for logistic regression.
 */
#include <mlpack/core.hpp>
#include "logistic_regression.hpp"

#include <mlpack/core/optimizers/sgd/sgd.hpp>

using namespace std;
using namespace mlpack;
using namespace mlpack::regression;
using namespace mlpack::optimization;

PROGRAM_INFO("L2-regularized Logistic Regression and Prediction",
    "An implementation of L2-regularized logistic regression using either the "
    "L-BFGS optimizer or SGD (stochastic gradient descent).  This solves the "
    "regression problem"
    "\n\n"
    "  y = (1 / 1 + e^-(X * b))"
    "\n\n"
    "where y takes values 0 or 1."
    "\n\n"
    "This program allows loading a logistic regression model from a file (-i) "
    "or training a logistic regression model given training data (-t), or both "
    "those things at once.  In addition, this program allows classification on "
    "a test dataset (-T) and will save the classification results to the given "
    "output file (-o).  The logistic regression model itself may be saved with "
    "a file specified using the -m option."
    "\n\n"
    "The training data given with the -t option should have class labels as its"
    " last dimension (so, if the training data is in CSV format, labels should "
    "be the last column).  Alternately, the -l (--labels_file) option may be "
    "used to specify a separate file of labels."
    "\n\n"
    "When a model is being trained, there are many options.  L2 regularization "
    "(to prevent overfitting) can be specified with the -l option, and the "
    "optimizer used to train the model can be specified with the --optimizer "
    "option.  Available options are 'sgd' (stochastic gradient descent) and "
    "'lbfgs' (the L-BFGS optimizer).  There are also various parameters for the"
    " optimizer; the --max_iterations parameter specifies the maximum number of"
    " allowed iterations, and the --tolerance (-e) parameter specifies the "
    "tolerance for convergence.  For the SGD optimizer, the --step_size "
    "parameter controls the step size taken at each iteration by the optimizer."
    "  If the objective function for your data is oscillating between Inf and "
    "0, the step size is probably too large.  There are more parameters for the"
    " SGD and L-BFGS optimizers, but the C++ interface must be used to access "
    "these."
    "\n\n"
    "Optionally, the model can be used to predict the responses for another "
    "matrix of data points, if --test_file is specified.  The --test_file "
    "option can be specified without --input_file, so long as an existing "
    "logistic regression model is given with --model_file.  The output "
    "predictions from the logistic regression model are stored in the file "
    "given with --output_predictions."
    "\n\n"
    "This implementation of logistic regression does not support the general "
    "multi-class case but instead only the two-class case.  Any responses must "
    "be either 0 or 1.");

// Training parameters.
PARAM_STRING("training_file", "A file containing the training set (the matrix "
    "of predictors, X).", "t", "");
PARAM_STRING("labels_file", "A file containing labels (0 or 1) for the points "
    "in the training set (y).", "l", "");

// Optimizer parameters.
PARAM_DOUBLE("lambda", "L2-regularization parameter for training.", "L", 0.0);
PARAM_STRING("optimizer", "Optimizer to use for training ('lbfgs' or 'sgd').",
    "O", "lbfgs");
PARAM_DOUBLE("tolerance", "Convergence tolerance for optimizer.", "e", 1e-10);
PARAM_INT("max_iterations", "Maximum iterations for optimizer (0 indicates no "
    "limit).", "M", 10000);
PARAM_DOUBLE("step_size", "Step size for SGD optimizer.", "s", 0.01);

// Model loading/saving.
PARAM_STRING("input_model", "File containing existing model (parameters).", "i",
    "");
PARAM_STRING("output_model", "File to save trained logistic regression model "
    "to.", "m", "");

// Testing.
PARAM_STRING("test_file", "File containing test dataset.", "T", "");
PARAM_STRING("output_file", "If --test_file is specified, this file is "
    "where the predicted responses will be saved.", "o", "");
PARAM_DOUBLE("decision_boundary", "Decision boundary for prediction; if the "
    "logistic function for a point is less than the boundary, the class is "
    "taken to be 0; otherwise, the class is 1.", "d", 0.5);

int main(int argc, char** argv)
{
  CLI::ParseCommandLine(argc, argv);

  // Collect command-line options.
  const string trainingFile = CLI::GetParam<string>("training_file");
  const string labelsFile = CLI::GetParam<string>("labels_file");
  const double lambda = CLI::GetParam<double>("lambda");
  const string optimizerType = CLI::GetParam<string>("optimizer");
  const double tolerance = CLI::GetParam<double>("tolerance");
  const double stepSize = CLI::GetParam<double>("step_size");
  const size_t maxIterations = (size_t) CLI::GetParam<int>("max_iterations");
  const string inputModelFile = CLI::GetParam<string>("input_model");
  const string outputModelFile = CLI::GetParam<string>("output_model");
  const string testFile = CLI::GetParam<string>("test_file");
  const string outputFile = CLI::GetParam<string>("output_file");
  const double decisionBoundary = CLI::GetParam<double>("decision_boundary");

  // One of inputFile and modelFile must be specified.
  if (trainingFile.empty() && inputModelFile.empty())
    Log::Fatal << "One of --input_model or --training_file must be specified."
        << endl;

  // If no output file is given, the user should know that the model will not be
  // saved, but only if a model is being trained.
  if (outputFile.empty() && !trainingFile.empty())
    Log::Warn << "--output_model not given; trained model will not be saved."
        << endl;

  // Tolerance needs to be positive.
  if (tolerance < 0.0)
    Log::Fatal << "Tolerance must be positive (received " << tolerance << ")."
        << endl;

  // Optimizer has to be L-BFGS or SGD.
  if (optimizerType != "lbfgs" && optimizerType != "sgd")
    Log::Fatal << "--optimizer must be 'lbfgs' or 'sgd'." << endl;

  // Lambda must be positive.
  if (lambda < 0.0)
    Log::Fatal << "L2-regularization parameter (--lambda) must be positive ("
        << "received " << lambda << ")." << endl;

  // Decision boundary must be between 0 and 1.
  if (decisionBoundary < 0.0 || decisionBoundary > 1.0)
    Log::Fatal << "Decision boundary (--decision_boundary) must be between 0.0 "
        << "and 1.0 (received " << decisionBoundary << ")." << endl;

  if ((stepSize < 0.0) && (optimizerType == "sgd"))
    Log::Fatal << "Step size (--step_size) must be positive (received "
        << stepSize << ")." << endl;

  if (CLI::HasParam("step_size") && optimizerType == "lbfgs")
    Log::Warn << "Step size (--step_size) ignored because 'sgd' optimizer is "
        << "not being used." << endl;

  // These are the matrices we might use.
  arma::mat regressors;
  arma::Mat<size_t> responsesMat;
  arma::Row<size_t> responses;
  arma::mat testSet;
  arma::Row<size_t> predictions;

  // Load data matrix.
  if (!trainingFile.empty())
    data::Load(trainingFile, regressors, true);

  // Load the model, if necessary.
  LogisticRegression<> model(0, 0); // Empty model.
  if (!inputModelFile.empty())
  {
    data::Load(inputModelFile, "logistic_regression_model", model);
  }
  else
  {
    // Set the size of the parameters vector, if necessary.
    if (labelsFile.empty())
      model.Parameters() = arma::zeros<arma::vec>(regressors.n_rows - 1);
    else
      model.Parameters() = arma::zeros<arma::vec>(regressors.n_rows);
  }

  // Check if the responses are in a separate file.
  if (!trainingFile.empty() && !labelsFile.empty())
  {
    data::Load(labelsFile, responsesMat, true);
    if (responsesMat.n_cols == 1)
      responses = responsesMat.col(0).t();
    else
      responses = responsesMat.row(0);

    if (responses.n_cols != regressors.n_cols)
      Log::Fatal << "The labels (--labels_file) must have the same number of "
          << "points as the training dataset (--training_file)." << endl;
  }
  else if (!trainingFile.empty())
  {
    // The initial predictors for y, Nx1.
    responses = arma::conv_to<arma::Row<size_t>>::from(
        regressors.row(regressors.n_rows - 1));
    regressors.shed_row(regressors.n_rows - 1);
  }

  // Verify the labels.
  if (!trainingFile.empty() && max(responses) > 1)
    Log::Fatal << "The labels must be either 0 or 1, not " << max(responses)
        << "!" << endl;

  // Now, do the training.
  if (!trainingFile.empty())
  {
    LogisticRegressionFunction<> lrf(regressors, responses, model.Parameters());
    if (optimizerType == "sgd")
    {
      SGD<LogisticRegressionFunction<>> sgdOpt(lrf);
      sgdOpt.MaxIterations() = maxIterations;
      sgdOpt.Tolerance() = tolerance;
      sgdOpt.StepSize() = stepSize;
      Log::Info << "Training model with SGD optimizer." << endl;

      // This will train the model.
      model.Train(sgdOpt);
    }
    else if (optimizerType == "lbfgs")
    {
      L_BFGS<LogisticRegressionFunction<>> lbfgsOpt(lrf);
      lbfgsOpt.MaxIterations() = maxIterations;
      lbfgsOpt.MinGradientNorm() = tolerance;
      Log::Info << "Training model with L-BFGS optimizer." << endl;

      // This will train the model.
      model.Train(lbfgsOpt);
    }
  }

  if (!testFile.empty())
  {
    data::Load(testFile, testSet, true);

    // We must perform predictions on the test set.  Training (and the
    // optimizer) are irrelevant here; we'll pass in the model we have.
    Log::Info << "Predicting classes of points in '" << testFile << "'."
        << endl;
    model.Predict(testSet, predictions, decisionBoundary);

    // Save the results, if necessary.
    if (!outputFile.empty())
      data::Save(outputFile, predictions, false);
  }

  if (!outputModelFile.empty())
  {
    Log::Info << "Saving model to '" << outputModelFile << "'." << endl;
    data::Save(outputModelFile, "logistic_regression_model", model, false);
  }
}
